package com.ppch.wuwamanus.rag;

import com.alibaba.cloud.ai.dashscope.agent.DashScopeAgentRagOptions;
import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetrieverOptions;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;

/**
 * Package:ppch-ai-agent
 * ClassName:MyContextualQueryAugmenter
 *
 * @Author 泡泡茶壶
 * @Create 2025/7/16 11:10
 * @Version 1.0
 * Description:
 *  自定义上下文查询增强器
 */
public class MyContextualQueryAugmenter {

    /**
     * 自定义上下文查询增强器
     * @param emptyContextPromptTemplate 未查询到问题相关文档时的提示词
     * @param allowEmptyContext 是否允许空上下文,当未查询到问题相关文档时，如果设置为
     *                          true：则使用原始的用户查询词且使用大模型的原始能力回答该问题，
     *                          false：则使用自定义的空上下文提示词模板进行回答
     * @return ContextualQueryAugmenter 上下文查询增强器
     */
    public ContextualQueryAugmenter createContextualQueryAugmenter(String emptyContextPromptTemplate,Boolean allowEmptyContext){
        //提示词
        PromptTemplate promptTemplate = new PromptTemplate(emptyContextPromptTemplate);
        //上下文查询增强器：自定义未查询到问题相关文档时的提示词，是否允许空上下文
        ContextualQueryAugmenter contextualQueryAugmenter = ContextualQueryAugmenter.builder()
                //当未查询到问题相关文档时的提示词模板
                .emptyContextPromptTemplate(promptTemplate)
                //是否允许空上下文，当未查询到问题相关文档时如果设置为true：则使用原始的用户查询词且使用大模型的原始能力回答该问题，
                //false：则使用自定义的空上下文提示词模板进行回答
                .allowEmptyContext(allowEmptyContext)
                .build();
        return contextualQueryAugmenter;
    }
}
